from langchain_core.documents import Document
from src.tools.mml_to_latex import mathml2latex
from uuid import uuid4

def to_document_bpr(file_dict, file):
    docus_para = []
    para_ids = []

    for para in file_dict["paragraphInfos"]:
        if para["text"] != "":
            para_id = str(uuid4())
            para_ids.append(para_id)
            docu_para = Document(page_content=para["text"], metadata={
                                                                "id": para_id,
                                                                 "index": para["index"],
                                                                 "sectionIndex": para["sectionIndex"],
                                                                 "bodyChildIndex": para["bodyChildIndex"],
                                                                 "text": para["text"]
                                                                 }
                            )
            docus_para.append(docu_para)

    docus_cell = []
    cell_ids = []
    for table in file_dict["tableInfos"]:
        for cell in table["tableCellInfos"]:
            cell_id = str(uuid4())
            cell_ids.append(cell_id)
            table_dict = {
                "id": cell_id,
                "index": table["index"],
                "sectionIndex": table["sectionIndex"],
                "bodyChildIndex": table["bodyChildIndex"],
                "tableCellInfos": []
            }
            cell_dict = {
                "rowIndex": cell["rowIndex"],
                "columnIndex": cell["columnIndex"],
                "paragraphInfos": []
            }
            cell_text = ""
            for para in cell["paragraphInfos"]:
                para_text = ""
                for textinfo in para["textInfos"]:
                    if textinfo["subSuperscript"] == 2:
                        para_text += f'<sub>{textinfo["text"]}</sub>'
                    else:
                        para_text += textinfo["text"]

                mathinfos = []
                for mathinfo in para["officeMathInfos"]:
                    math_dict = {
                        "text": mathml2latex(mathinfo["mathMlCode"])
                    }
                    mathinfos.append(math_dict)

                para_dict = {
                    "index": para["index"],
                    "sectionIndex": para["sectionIndex"],
                    "bodyChildIndex": para["bodyChildIndex"],
                    "text": para_text,
                    "latexMathInfos": mathinfos
                }
                cell_dict["paragraphInfos"].append(para_dict)
                cell_text += para_text
            table_dict["tableCellInfos"].append(cell_dict)
            docu_cell = Document(page_content=cell_text, metadata=table_dict)
            docus_cell.append(docu_cell)
    return docus_para, para_ids, docus_cell, cell_ids


def to_document_sop(file_dict, file):
    docus_para = []

    for para in file_dict["paragraphInfos"]:
        if para["text"] != "":
            docu_para = Document(page_content=para["text"], metadata={
                                                                 "index": para["index"],
                                                                 "sectionIndex": para["sectionIndex"],
                                                                 "bodyChildIndex": para["bodyChildIndex"],
                                                                 "text": para["text"]
                                                                 }
                            )
            docus_para.append(docu_para)

    docus_table = []

    for table in file_dict["tableInfos"]:
        table_dict = {
            "index": table["index"],
            "sectionIndex": table["sectionIndex"],
            "bodyChildIndex": table["bodyChildIndex"],
            "tableCellInfos": []
        }
        table_text = ""
        for cell in table["tableCellInfos"]:
            cell_dict = {
                "rowIndex": cell["rowIndex"],
                "columnIndex": cell["columnIndex"],
                "paragraphInfos": []
            }
            cell_text = ""
            for para in cell["paragraphInfos"]:
                para_text = ""
                for textinfo in para["textInfos"]:
                    if textinfo["subSuperscript"] == 2:
                        para_text += f'<sub>{textinfo["text"]}</sub>'
                    else:
                        para_text += textinfo["text"]

                mathinfos = []
                for mathinfo in para["officeMathInfos"]:
                    math_dict = {
                        "text": mathml2latex(mathinfo["mathMlCode"])
                    }
                    mathinfos.append(math_dict)

                para_dict = {
                    "index": para["index"],
                    "sectionIndex": para["sectionIndex"],
                    "bodyChildIndex": para["bodyChildIndex"],
                    "text": para_text,
                    "latexMathInfos": mathinfos
                }
                cell_dict["paragraphInfos"].append(para_dict)
                cell_text += para["text"]
            table_dict["tableCellInfos"].append(cell_dict)
            table_text += cell_text + "|"
        docu_table = Document(page_content=table_text, metadata=table_dict)
        docus_table.append(docu_table)
    return docus_para, docus_table


